#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
LLM客户端

通用的大模型对话客户端，支持命令行参数输入指令
"""

import json
import argparse
import sys
from openai import OpenAI


class LLMClient:
    """
    大模型客户端
    用于与大模型进行对话交互
    """
    
    def __init__(self):
        # 大模型配置
        self.model_default_headers = {
            "Content-Type": "application/json",
            "source-sn": "ai-math-class-agent"  # 替换为审批通过的服务标识
        }

        self.model_name = "doubao-seed-1-6-250615"
        
        try:
            # 初始化OpenAI客户端
            self.client = OpenAI(
                api_key="sk-boM5kzJgrXVFRuZ85f9fE3568372484aA688F1Ad9f9257Ec",
                base_url="http://oneapi.dev.service.yuaiweiwu.com/v1",
                default_headers=self.model_default_headers,
            )
            
            # 测试连接
            self._test_connection()
            print(" 大模型服务连接成功!")
            
        except Exception as e:
            print(f" 大模型服务连接失败: {e}")
            sys.exit(1)
    
    def _test_connection(self):
        """测试大模型连接"""
        try:
            response = self.client.chat.completions.create(
                model=self.model_name,
                messages=[
                    {"role": "system", "content": "你是一个智能助手。"},
                    {"role": "user", "content": "测试连接"}
                ],
                max_tokens=10,
                timeout=10
            )
            return True
        except Exception as e:
            raise Exception(f"连接测试失败: {e}")
    
    def chat(self, user_input, system_prompt=None, max_tokens=250, temperature=0.7):
        """
        与大模型进行对话
        
        Args:
            user_input: 用户输入的指令或问题
            system_prompt: 系统提示词（可选）
            max_tokens: 最大token数（默认250，适合语音播报）
            temperature: 温度参数
            
        Returns:
            大模型的回复内容
        """
        try:
            messages = []
            
            # 添加系统提示词
            if system_prompt:
                messages.append({"role": "system", "content": system_prompt})
            else:
                messages.append({"role": "system", "content": "你是一个智能助手，请根据用户的问题提供准确、有用的回答。"})
            
            # 添加用户输入
            messages.append({"role": "user", "content": user_input})
            
            print(f" 正在处理您的请求...")
            
            response = self.client.chat.completions.create(
                model=self.model_name,
                messages=messages,
                max_tokens=max_tokens,
                temperature=temperature,
                timeout=120  # 增加超时时间到120秒
            )
            
            result = response.choices[0].message.content.strip()
            return result
            
        except Exception as e:
            print(f" 大模型调用失败: {e}")
            return None
    
    def generate_friendly_error_response(self, user_command, error_reason):
        """
        生成友好的错误回复
        
        Args:
            user_command: 用户的原始指令
            error_reason: 错误原因
            
        Returns:
            友好的错误回复
        """
        try:
            system_prompt = """你是一个智能且友好的AI助手。当用户的指令无法被系统执行时，你需要：

**判断用户意图：**
1. 如果用户是想让你生成内容（如写小说、写文章、写代码、翻译、回答问题等），请直接满足用户的需求，生成相应内容
2. 如果用户是想执行系统操作（如播放音乐、打开应用、控制设备等），这类指令确实无法执行，才需要温和地说明原因并给出建议

**对于内容生成类请求：**
- 写小说/文章 → 直接写一篇简短的内容（⚠️ 严格控制在200字以内，因为需要语音播报）
- 写代码 → 直接提供简短代码示例（不超过20行）
- 翻译 → 直接进行翻译
- 回答问题 → 直接回答（简洁，不超过150字）
不要说"无法执行"或"错误"

**对于系统操作类请求：**
- 用温和友好的语气说明原因
- 给出1-2个可用的指令示例
- 保持回复简洁（50字以内）

注意：不要使用emoji表情，使用纯文字回复。内容要精炼，适合语音播报。"""
            
            user_prompt = f"""用户说了："{user_command}"
系统返回：{error_reason}

请判断用户意图并给出合适的回复：
- 如果是内容生成请求，直接生成内容（⚠️ 严格控制在200字以内）
- 如果是系统操作请求，说明原因并给出建议"""
            
            response = self.chat(
                user_input=user_prompt,
                system_prompt=system_prompt,
                max_tokens=250,  # 减少到250，严格控制生成长度
                temperature=0.8  # 提高温度，让生成的内容更有创意
            )
            
            return response
            
        except Exception as e:
            # 如果LLM调用失败，返回默认友好回复
            print(f"⚠️ 生成友好回复失败，使用默认回复: {e}")
            return f"抱歉，我无法理解您的指令「{user_command}」。您可以试试说：播放音乐、打开记事本、创建文件等。"
    
    def generate_friendly_error_response_stream(self, user_command, error_reason):
        """
        流式生成友好的错误回复
        
        Args:
            user_command: 用户的原始指令
            error_reason: 错误原因
            
        Yields:
            流式生成的文本片段
        """
        try:
            system_prompt = """你是一个智能且友好的AI助手。当用户的指令无法被系统执行时，你需要：

**判断用户意图：**
1. 如果用户是想让你生成内容（如写小说、写文章、写代码、翻译、回答问题等），请直接满足用户的需求，生成相应内容
2. 如果用户是想执行系统操作（如播放音乐、打开应用、控制设备等），这类指令确实无法执行，才需要温和地说明原因并给出建议

**对于内容生成类请求：**
- 写小说/文章 → 直接写一篇简短的内容（⚠️ 严格控制在200字以内，因为需要语音播报）
- 写代码 → 直接提供简短代码示例（不超过20行）
- 翻译 → 直接进行翻译
- 回答问题 → 直接回答（简洁，不超过150字）
不要说"无法执行"或"错误"

**对于系统操作类请求：**
- 用温和友好的语气说明原因
- 给出1-2个可用的指令示例
- 保持回复简洁（50字以内）

注意：不要使用emoji表情，使用纯文字回复。内容要精炼，适合语音播报。"""
            
            user_prompt = f"""用户说了："{user_command}"
系统返回：{error_reason}

请判断用户意图并给出合适的回复：
- 如果是内容生成请求，直接生成内容（⚠️ 严格控制在200字以内）
- 如果是系统操作请求，说明原因并给出建议"""
            
            messages = [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt}
            ]
            
            print(f"🌊 开始流式生成智能回复...")
            
            # 使用流式API - 严格控制生成长度，避免TTS超限
            stream = self.client.chat.completions.create(
                model=self.model_name,
                messages=messages,
                max_tokens=250,  # 减少到250，严格控制生成长度
                temperature=0.8,  # 提高温度，让生成的内容更有创意
                stream=True,
                timeout=120
            )
            
            chunk_count = 0
            has_generated_content = False  # 标记是否已生成内容
            
            for chunk in stream:
                # 检查 chunk 是否有 choices，并且 choices 不为空
                if chunk.choices and len(chunk.choices) > 0:
                    if chunk.choices[0].delta.content:
                        content = chunk.choices[0].delta.content
                        chunk_count += 1
                        has_generated_content = True  # 标记已有内容
                        print(f"📦 [Chunk {chunk_count}] {content}")
                        yield content
            
            print(f"✅ 流式生成完成，共 {chunk_count} 个片段")
            
        except Exception as e:
            print(f"⚠️ 流式生成出现异常: {e}")
            
            # 只有在完全没有生成任何内容时，才返回默认回复
            if not has_generated_content:
                print(f"💬 完全没有生成内容，使用默认回复")
                default_message = f"抱歉，我无法理解您的指令「{user_command}」。您可以试试说：播放音乐、创建文件等。"
                yield default_message
            else:
                print(f"💬 已生成 {chunk_count} 个片段，流式中断但不添加默认消息")
                # 不 yield 任何内容，让已生成的内容保持原样
                # 异常会在 app.py 中的 except 块被捕获和处理
    
    def chat_stream(self, user_input, system_prompt=None, max_tokens=250, temperature=0.7):
        """
        流式对话
        
        Args:
            user_input: 用户输入的指令或问题
            system_prompt: 系统提示词（可选）
            max_tokens: 最大token数（默认250，适合语音播报）
            temperature: 温度参数
            
        Yields:
            流式生成的文本片段
        """
        chunk_count = 0
        has_generated_content = False
        
        try:
            messages = []
            
            # 添加系统提示词
            if system_prompt:
                messages.append({"role": "system", "content": system_prompt})
            else:
                messages.append({"role": "system", "content": "你是一个智能且友好的AI助手，请用简洁、自然的方式回答用户的问题。"})
            
            # 添加用户输入
            messages.append({"role": "user", "content": user_input})
            
            print(f"🌊 开始流式生成回复...")
            
            # 使用流式API
            stream = self.client.chat.completions.create(
                model=self.model_name,
                messages=messages,
                max_tokens=max_tokens,
                temperature=temperature,
                stream=True,
                timeout=120
            )
            
            for chunk in stream:
                # 检查 chunk 是否有 choices，并且 choices 不为空
                if chunk.choices and len(chunk.choices) > 0:
                    if chunk.choices[0].delta.content:
                        content = chunk.choices[0].delta.content
                        chunk_count += 1
                        has_generated_content = True
                        yield content
            
            print(f"✅ 流式对话完成，共 {chunk_count} 个片段")
            
        except Exception as e:
            print(f"❌ 流式对话出现异常: {type(e).__name__}: {e}")
            import traceback
            traceback.print_exc()
            
            # 只有在完全没有生成任何内容时，才返回默认回复
            if not has_generated_content:
                print(f"💬 完全没有生成内容，返回默认回复")
                yield "抱歉，我现在无法回答您的问题。请稍后再试。"
            else:
                print(f"💬 已生成 {chunk_count} 个片段，流式中断但不添加默认消息")
    
    def chat_with_context(self, messages, max_tokens=1000, temperature=0.7):
        """
        带上下文的多轮对话
        
        Args:
            messages: 消息列表，格式为 [{"role": "user/assistant/system", "content": "内容"}]
            max_tokens: 最大token数
            temperature: 温度参数
            
        Returns:
            大模型的回复内容
        """
        try:
            print(f" 正在处理多轮对话...")
            
            response = self.client.chat.completions.create(
                model=self.model_name,
                messages=messages,
                max_tokens=max_tokens,
                temperature=temperature,
                timeout=120  # 增加超时时间到120秒
            )
            
            result = response.choices[0].message.content.strip()
            return result
            
        except Exception as e:
            print(f" 大模型调用失败: {e}")
            return None


def main():
    """命令行主函数"""
    parser = argparse.ArgumentParser(description='LLM客户端 - 与大模型进行对话')
    parser.add_argument('prompt', nargs='?', help='要发送给大模型的指令或问题')
    parser.add_argument('-s', '--system', help='系统提示词')
    parser.add_argument('-t', '--temperature', type=float, default=0.7, help='温度参数 (0.0-1.0)')
    parser.add_argument('-m', '--max-tokens', type=int, default=1000, help='最大token数')
    parser.add_argument('-i', '--interactive', action='store_true', help='进入交互模式')
    parser.add_argument('--json', action='store_true', help='以JSON格式输出结果')
    
    args = parser.parse_args()
    
    # 初始化客户端
    try:
        client = LLMClient()
    except Exception as e:
        print(f" 客户端初始化失败: {e}")
        sys.exit(1)
    
    # 交互模式
    if args.interactive:
        print(" 进入交互模式，输入 'quit' 或 'exit' 退出")
        print("=" * 50)
        
        conversation_history = []
        if args.system:
            conversation_history.append({"role": "system", "content": args.system})
        
        while True:
            try:
                user_input = input("\n 您: ").strip()
                
                if user_input.lower() in ['quit', 'exit', '退出', 'q']:
                    print(" 再见！")
                    break
                
                if not user_input:
                    continue
                
                # 添加用户消息到历史
                conversation_history.append({"role": "user", "content": user_input})
                
                # 获取回复
                response = client.chat_with_context(
                    conversation_history, 
                    max_tokens=args.max_tokens, 
                    temperature=args.temperature
                )
                
                if response:
                    print(f"\n 助手: {response}")
                    # 添加助手回复到历史
                    conversation_history.append({"role": "assistant", "content": response})
                else:
                    print(" 获取回复失败")
                
            except KeyboardInterrupt:
                print("\n 再见！")
                break
            except Exception as e:
                print(f" 处理输入时出错: {e}")
    
    # 单次对话模式
    elif args.prompt:
        response = client.chat(
            args.prompt, 
            system_prompt=args.system,
            max_tokens=args.max_tokens,
            temperature=args.temperature
        )
        
        if response:
            if args.json:
                result = {
                    "success": True,
                    "prompt": args.prompt,
                    "response": response,
                    "model": client.model_name
                }
                print(json.dumps(result, ensure_ascii=False, indent=2))
            else:
                print(f"\n 回复: {response}")
        else:
            if args.json:
                result = {
                    "success": False,
                    "prompt": args.prompt,
                    "error": "获取回复失败"
                }
                print(json.dumps(result, ensure_ascii=False, indent=2))
            else:
                print(" 获取回复失败")
            sys.exit(1)
    
    # 没有提供参数，显示帮助
    else:
        parser.print_help()
        print("\n示例用法:")
        print("  python llm_client.py '你好，请介绍一下自己'")
        print("  python llm_client.py '写一个Python函数计算斐波那契数列' -s '你是一个Python编程专家'")
        print("  python llm_client.py -i  # 进入交互模式")
        print("  python llm_client.py '解释量子计算' --json  # JSON格式输出")


if __name__ == "__main__":
    main()
